﻿using System;
using System.Collections.Generic;
using System.Linq;
using Algoritmia.StateSpaces;
using Algoritmia.Utils;
using Lists = Algoritmia.DataStructures.Lists;

namespace Algoritmia.Problems.Puzzles.NQueens
{

    public class NQueensStateSpace2 : IReversibleForwardStateSpace<NQueensStateSpace2.State,int>
    {
        private readonly int _n;

        public class State
        {
            public readonly State Prev;
            public readonly int Row;
            public readonly int N;
            public List<int> Rows { get; private set; }

            public State(State prev = null, int row = 0, int n = 0)
            {
                Prev = prev;
                Row = row;
                N = n;
                Rows = prev != null ? prev.Rows : new List<int>();
            }       

            public Lists.IList<int> ToList()
            {
                State s = this;
                int[] array = new int[N];
                while (s.Prev != null)
                {
                    array[s.N - 1] = s.Row;
                    s = s.Prev;
                }
                return AlgoritmiaFactory.Lists<int>.BuildDefault(array);
            }

        }

        public NQueensStateSpace2(int N)
        {
            _n = N;
        }

        public State Decide(State state, int decision)
        {
            var newState = new State(state, decision, state.N + 1);
            newState.Rows.Add(decision);
            return newState;
        }

        public IEnumerable<int> Decisions(State state)
        {
            for(int row = 0; row < _n; row++)
            {
                if ( (state.N == 0) || (!state.Rows.Contains(row) && Enumerable.Range(0,state.N).All( j => state.N - j != Math.Abs(row-state.Rows[j]))))
                {
                    yield return row;
                }
            }
        }

        public IEnumerable<State> InitialStates
        {
            get { yield return new State(); }
        }

        public bool IsFinal(State state)
        {
            return _n == state.N;
        }

        public State Undo(State state, int decision)
        {
            state.Rows.RemoveAt(state.Rows.Count-1);
            return state.Prev;
        }
    }


    public class NQueensEnumerator2
    {
        protected readonly Func<Lists.IList<int>> _createList;

        public NQueensEnumerator2(Func<Lists.IList<int>> createList = null)
        {
            _createList = createList ?? ( () => AlgoritmiaFactory.Lists<int>.Default ) ;
        }

        protected virtual IEnumerable<NQueensStateSpace2.State> Backtracking(NQueensStateSpace2.State state, 
            NQueensStateSpace2 space, Lists.IList<int> decisions)
        {
            if (space.IsFinal(state))
            {
                yield return state;
            }

            foreach (var decision in space.Decisions(state))
            {
                decisions.Add(decision);
                state = space.Decide(state, decision);
                foreach (var result in Backtracking(state, space, decisions))
                {
                    yield return result;
                }
                decisions.RemoveAt(decisions.Count-1);
                state = space.Undo(state, decision);
            }
        }

        public IEnumerable<Lists.IList<int>> Enumerate(NQueensStateSpace2 space)
        {
            var decisions = _createList();
            var initial = space.InitialStates.First();
            if (initial != null)
            {
                foreach (var result in Backtracking(initial, space, decisions))
                {
                    yield return result.ToList();
                }
            }
        }

        public Lists.IList<int> FirstSolution(NQueensStateSpace2 space)
        {
            return Enumerate(space).First();
        }

    }
}
